import cv2
import numpy as np
import h5py


def list_to_npy_trainval(list_path, npy_path_im, npy_path_la, resize):
    train_folder_path = '/media/dell/delldisk/dell/wxm/Data/KaggleDDD/train'
    f_list = open(list_path, 'r')
    img_arrs = []
    labels = []
    for line in f_list:
        sub_path, label = line.split(' ')
        img_path = train_folder_path + '/' + sub_path
        img_arr = cv2.imread(img_path)
        img_arr = cv2.cvtColor(img_arr, cv2.COLOR_RGB2BGR)
        img_arr = img_arr[80:400, 320:]
        img_arr = img_arr[:, 80:560]
        img_arr = cv2.resize(img_arr, (resize, resize), interpolation=cv2.INTER_AREA)
        img_arrs.append(img_arr)
        labels.append(label)
    img_arrs = np.asarray(img_arrs, np.uint8)
    labels = np.asarray(labels, np.uint8)
    np.save(npy_path_im, img_arrs)
    np.save(npy_path_la, labels)

def list_to_h5_trainval(list_path, h5_path, resize):
    train_folder_path = '/media/dell/delldisk/dell/wxm/Data/KaggleDDD/train'
    f_list = open(list_path, 'r')
    img_arrs = []
    labels = []
    for line in f_list:
        sub_path, label = line.split(' ')
        img_path = train_folder_path + '/' + sub_path
        img_arr = cv2.imread(img_path)
        # img_arr = cv2.cvtColor(img_arr, cv2.COLOR_RGB2BGR)
        img_arr = img_arr[80:400, 320:]
        img_arr = img_arr[:, 80:560]
        img_arr = cv2.resize(img_arr, (resize, resize), interpolation=cv2.INTER_AREA)
        img_arrs.append(img_arr)
        labels.append(label)
    f = h5py.File(h5_path)
    img_arrs = np.asarray(img_arrs, np.uint8)
    labels = np.asarray(labels, np.uint8)
    f.create_dataset('data', data=img_arrs)
    f.create_dataset('label', data=labels)
    f.close()


list_to_h5_trainval(
    list_path='/media/dell/delldisk/dell/wxm/Data/KaggleDDD/experiments/caffebased/data/val_list.txt',
    h5_path='/media/dell/delldisk/dell/wxm/Data/KaggleDDD/experiments/caffebased/data/val_caffe_RGB.h5',
    resize=256)
